# -*- coding: utf-8 -*-
import random
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
import itertools


class Conv:
    cs_padding = CS.ConfigurationSpace()
    cs_attributes = CS.ConfigurationSpace()
    attribute_list = ['auto_pad', 'dilations', 'group', 'kernel_shape', 'pads', 'strides']
    auto_pad_list = ['NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID']
    all_attributes = []

    def __init__(self):
        print('Conv Init--------------------------')
        self.init_all_attributes()
        self.init_cs()

    def init_all_attributes(self):
        for i in range(1, len(Conv.attribute_list)+1) :
            #此处取所有组合的可能值，存在列表里
            for attribute in itertools.combinations(Conv.attribute_list, i):
                l = list(attribute)
                Conv.all_attributes.append(l)

        #print('all_attributes: ', all_attributes)
        print('all attributes:')
        for a in Conv.all_attributes:
            print(a)

        tmp_attributes = []

        #将所有含有auto_pad的元素展开为'NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID'四种类型
        #如旧元素为['group', 'auto_pad']，则展开后的元素为：
        #['group', 'NOTSET']
        #['group', 'SAME_UPPER']
        #['group', 'SAME_LOWER']
        #['group', 'VALID']

        for attribute in Conv.all_attributes :
            tmp_a = []

            if 'auto_pad' in attribute : 
                tmp = []
                for a in attribute :
                    if a != 'auto_pad' :
                        tmp.append(a)

                tmp1 = tmp.copy()
                tmp2 = tmp.copy() 
                tmp3 = tmp.copy() 
                tmp4 = tmp.copy() 

                #print('tmp1:', tmp1)
                #print('tmp2:', tmp2)
                #print('tmp3:', tmp3)
                #print('tmp4:', tmp4)

                tmp1.append('NOTSET') 
                tmp2.append('SAME_UPPER')
                tmp3.append('SAME_LOWER')       
                tmp4.append('VALID')

                #print('--------tmp1:', tmp1)
                #print('--------tmp2:', tmp2)
                #print('--------tmp3:', tmp3)
                #print('--------tmp4:', tmp4)

                tmp_a.append(tmp1)
                tmp_a.append(tmp2)
                tmp_a.append(tmp3)
                tmp_a.append(tmp4)

                #print('+++++++++++++++++++++++++tmp:', tmp_a)

                tmp_attributes.append(tmp_a) 

        print('got new attributes: ')   
        for a in tmp_attributes:
                print(a) 

        for a in Conv.all_attributes[:]:
            if 'auto_pad' in a : 
                Conv.all_attributes.remove(a)

        #print('got last attributes: ') 
        for a in Conv.all_attributes:
            print(a)

        for a in tmp_attributes:
            Conv.all_attributes.extend(a)   

        print('again, got last attributes: ') 
        for a in Conv.all_attributes:
            print(a)         
                
    def init_cs(self):
        padding = CSH.CategoricalHyperparameter('auto_pad', ['NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID'])
        Conv.cs_padding.add_hyperparameter(padding)

        length = len(Conv.all_attributes)

        attributes = CSH.UniformIntegerHyperparameter(name='combination', lower=0, upper=2*length)
        Conv.cs_attributes.add_hyperparameter(attributes)

    def get_random_padding(self):
        s = cs_padding.sample_configuration()
        padding = s.get_dictionary()['auto_pad']
        print('get auto_pad: ', padding)

    def get_random_combination(self):
        s = Conv.cs_attributes.sample_configuration()
        index = s.get_dictionary()['combination']
        print('get combination: ', index)

        length = len(Conv.all_attributes)

        #保证不带任何属性的case可以被覆盖到
        if index >= length:
            print('test no attributes')
            return []

        return Conv.all_attributes[index]            

    def input_shape_correct(self, shapes):
        print('input_shape_correct, Conv')

        if len(shapes) != 2 and len(shapes) != 3:
            return False

        if len(shapes) == 3:
            shapes[2] = [100]    

        if len(shapes[0]) > 4 :
            shapes[0] = shapes[0][:4]

        if len(shapes[1]) > 4 :
            shapes[1] = shapes[1][:4]   

        if len(shapes[0]) < 3 and len(shapes[1]) < 3 :
            if len(shapes[0]) % 2 == 0:
                #print('reshape input 111111111')
                for i in range(4-len(shapes[0])):
                    shapes[0].append(random.randint(1, 8))
            else:
                #print('reshape input 222222222222222')
                for i in range(3-len(shapes[0])):
                    shapes[0].append(random.randint(8, 16))

        if len(shapes[0]) != len(shapes[1]) : 
            if len(shapes[0]) < 3:
                shapes[0] = shapes[1].copy()
            elif len(shapes[1]) < 3:
                shapes[1] = shapes[0].copy()    
            else:
                if len(shapes[0]) % 2 == 0:
                    #print('change shapes[1]', shapes[1], shapes[0])
                    shapes[1] = shapes[0].copy()
                else:
                    #print('change shapes[0]', shapes[1], shapes[0])
                    shapes[0] = shapes[1].copy()                       

        #print('-----before change shapes[1][0]', shapes[1][0])

        shapes[1][0] = random.randint(shapes[0][0], 10*shapes[0][0])

        #print('-----after change shapes[1][0]', shapes[1][0], shapes[0][0])

        v = int(shapes[0][1]/shapes[1][1])
        if v != 0:
            shapes[0][1] = v * shapes[1][1]
        else :
            v = int(shapes[1][1]/shapes[0][1])
            if v != 0 :
                shapes[0][1], shapes[1][1] = shapes[1][1], shapes[0][1]
                shapes[0][1] = v * shapes[1][1]
            else:    
                shapes[0][1] = shapes[1][1]

        v =  int(shapes[0][1]/shapes[1][1])
        tmp = int(shapes[1][0]/v)
        if tmp == 0:
            tmp = v

        shapes[1][0] = tmp * v

        print('set group: ', shapes[0][1], shapes[1][1], v, tmp, shapes[1][0])          

        if len(shapes) == 3:
            #print('KKKKKKKKKKKKKKKKKKK')
            if shapes[2][0] > 1 :
                print('change shapes[2][0]', shapes[2][0], shapes[1][0])
                shapes[2][0] = shapes[1][0]

        #print('-----before change shapes[0][2]', shapes[0][2])

        shapes[0][2] = random.randint(shapes[1][2], 2*shapes[1][2])

        print('-----after change shapes[0][2]', shapes[0][2])

        #print('last length:', len(shapes[0]), len(shapes[1]))

        if shapes[1][2] > shapes[0][2] : 
            print('ERROR, should not go here~~~~', shapes[1][2], shapes[0][2])
            shapes[1][2] = shapes[0][2] 

        if len(shapes[0]) == 4:
            #print('-----before change shapes[0][3]', shapes[0][3])
            shapes[0][3] = random.randint(shapes[1][3], 2*shapes[1][3])
            #print('-----after change shapes[0][3]', shapes[0][3])

            if shapes[1][3] > shapes[0][3] : 
                print('XXXX ERROR, should not go here~~~~', shapes[1][2], shapes[0][2])
                shapes[1][3] = shapes[0][3]

        #for dim in enumerate(shapes[0]):
        #    print('last shapes[0]: ', dim)    

        #for dim in enumerate(shapes[1]):
        #    print('last shapes[1]: ', dim)        

    def infer_output_shape(self, shapes):
        shape = []

        if len(shapes) != 2 and len(shapes) != 3:
            return shape, False

        if len(shapes[0]) != 3 and len(shapes[0]) != 4:
            return shape, False

        if len(shapes[1]) != 3 and len(shapes[1]) != 4:
            return shape, False

        if len(shapes[0]) != len(shapes[1]) :
            return shape, False                

        if len(shapes[0]) == 4:
            shape.append(shapes[0][0])
            shape.append(shapes[1][0])
            shape.append(shapes[0][2] - shapes[1][2] + 1)
            shape.append(shapes[0][3] - shapes[1][3] + 1)
        elif len(shapes[0]) == 3:  
            shape.append(shapes[0][0])
            shape.append(shapes[1][0])
            shape.append(shapes[0][2] - shapes[1][2] + 1)
    
        return shape, True    

    def get_max_attributes_combination(self):
        #global _init_cs
        '''
        if Conv._init_cs == False :
            self.init_all_attributes()
            self.init_cs()
            Conv._init_cs = True
        '''    

        return len(Conv.all_attributes)

    def get_attributes_by_index_or_random(self, shapes, index):
        #global _init_cs

        '''
        if Conv._init_cs == False :
            init_all_attributes()
            init_cs()
            _init_cs = True
        '''
        print('get_attributes_by_index_or_random, index: ', index)    
            
        attributes = {}
        shape_out = []

        stride_ = [1, 1]
        dilation_ = [1, 1]
        pads_ = [0, 0, 0, 0]

        has_dilations = False
        has_pads = False
        has_strides = False

        auto_pad_mode = 'NOTSET'

        group = 1

        if index != -1:
            attr_combination = Conv.all_attributes[index]
        else:
            attr_combination = self.get_random_combination()

        if len(attr_combination) == 0:
            v = int(shapes[0][1]/shapes[1][1])
            if v != 1:
                attributes['group'] = v
                group = v
                
            return attributes, shape_out    

        if 'kernel_shape' in attr_combination:
            if len(shapes[1]) == 4:
                kernel_shape = [shapes[1][2], shapes[1][3]]
                attributes['kernel_shape'] = kernel_shape
            else:
                kernel_shape = [shapes[1][2]]
                attributes['kernel_shape'] = kernel_shape    

        if 'strides' in attr_combination:
            if len(shapes[0]) == 4:
                strides = [random.randint(1, shapes[0][2]+1), random.randint(1, shapes[0][3])]
                attributes['strides'] = strides
                stride_[0] = strides[0]
                stride_[1] = strides[1]
                has_strides = True
            else:
                strides = [random.randint(1, shapes[0][2]+1)]
                attributes['strides'] = strides
                stride_[0] = strides[0] 
                has_strides = True   

        if 'group' in attr_combination:
            group = int(shapes[0][1]/shapes[1][1]) #random.randint(1, shapes[0][1]+1)
            attributes['group'] = group
            #print('before adjust, ', shapes[0][1])
            #shapes[0][1] = group * shapes[0][1]

        if 'group' not in attr_combination:
            v = int(shapes[0][1]/shapes[1][1])
            if v != 1:
                attributes['group'] = v
                group = v

        has_auto_pad = False        

        for attr in attr_combination:
            if attr in Conv.auto_pad_list:
                index = Conv.auto_pad_list.index(attr)
                print('set auto_pad: ', Conv.auto_pad_list[index])
                attributes['auto_pad'] = Conv.auto_pad_list[index]
                auto_pad_mode = Conv.auto_pad_list[index]
                has_auto_pad = True
                print('set auto_pad_mode: ', auto_pad_mode)

        #dilations属性不能与SAME_UPPER或SAME_LOWER属性同时存在
        if auto_pad_mode != 'SAME_UPPER' and auto_pad_mode != 'SAME_LOWER' :
            if 'dilations' in attr_combination:
                if len(shapes[0]) == 4:
                    dilations = [random.randint(1, shapes[0][2]+1), random.randint(1, shapes[0][3])]
                    attributes['dilations'] = dilations
                    dilation_[0] = dilations[0]
                    dilation_[1] = dilations[1]
                    has_dilations = True
                else:
                    dilations = [random.randint(1, shapes[0][2]+1)]
                    attributes['dilations'] = dilations
                    dilation_[0] = dilations[0]
                    has_dilations = True        

        if has_auto_pad == False:
            if 'pads' in attr_combination : 
                if len(shapes[0]) == 4:
                    pads = [random.randint(0, int((shapes[0][2]-shapes[1][2])/2) + 1), random.randint(0, int((shapes[0][3]-shapes[1][3])/2) + 1),
                                random.randint(0, int((shapes[0][2]-shapes[1][2])/2) + 1), random.randint(0, int((shapes[0][3]-shapes[1][3])/2) + 1)]
                    attributes['pads'] = pads
                    pads_[0] = pads[0]
                    pads_[1] = pads[1]
                    pads_[2] = pads[2]
                    pads_[3] = pads[3]
                    has_pads = True
                else:
                    pads = [random.randint(0, int((shapes[0][2]-shapes[1][2])/2) + 1), random.randint(0, int((shapes[0][2]-shapes[1][2])/2) + 1)]
                    attributes['pads'] = pads
                    pads_[0] = pads[0]
                    pads_[1] = pads[1]
                    has_pads = True

        shape_out.append(shapes[0][0])
        shape_out.append(group*shapes[1][0])

        if len(shapes[0]) == 4:
            if int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pads_[0] + pads_[2])/stride_[0]) <= 0 :
                v = (shapes[1][2] - 1)*dilation_[0]  + 1 + stride_[0] - pads_[0] - pads_[2]
                if v > 0 :
                    #print('before reshape shapes[0][2]: ', shapes[0][2])
                    shapes[0][2] = random.randint(v, 2*v)
                    #print('after reshape shapes[0][2]: ', shapes[0][2])
                else:
                    shapes[0][2] = 1 
                    #print('reshape shapes[0][2] to 1') 

            if int((shapes[0][3] - ((shapes[1][3] - 1)*dilation_[1] + 1) + pads_[1] + pads_[3])/stride_[1]) <= 0 :
                v = (shapes[1][3] - 1)*dilation_[1]  + 1 + stride_[1] - pads_[1] - pads_[3]
                if v > 0 :
                    #print('before reshape shapes[0][3]: ', shapes[0][3])
                    shapes[0][3] = random.randint(v, 2*v)
                    #print('after reshape shapes[0][3]: ', shapes[0][3])
                else:
                    shapes[0][3] = 1 
                    #print('reshape shapes[0][3] to 1') 

        if len(shapes[0]) == 3:
            if int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pads_[0] + pads_[1])/stride_[0]) <= 0 :
                v = (shapes[1][2] - 1)*dilation_[0]  + 1 + stride_[0] - pads_[0] - pads_[1]
                if v > 0 :
                    #print('-----before reshape shapes[0][2]: ', shapes[0][2])
                    shapes[0][2] = random.randint(v, 2*v)
                    #print('-----after reshape shapes[0][2]: ', shapes[0][2])
                else:
                    shapes[0][2] = 1 
                    #print('-----reshape shapes[0][2] to 1')                         

        if len(shapes[0]) == 4:
            if auto_pad_mode == 'NOTSET' :
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pads_[0] + pads_[2])/stride_[0] + 1))
                shape_out.append(int((shapes[0][3] - ((shapes[1][3] - 1)*dilation_[1] + 1) + pads_[1] + pads_[3])/stride_[1] + 1)) 
            elif auto_pad_mode == 'VALID' :
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1))/stride_[0] + 1))
                shape_out.append(int((shapes[0][3] - ((shapes[1][3] - 1)*dilation_[1] + 1))/stride_[1] + 1)) 
            elif auto_pad_mode == 'SAME_UPPER' or auto_pad_mode == 'SAME_LOWER' : 
                legacy_target_size_h = int((shapes[0][2] + stride_[0] - 1) / stride_[0])
                legacy_target_size_w = int((shapes[0][3] + stride_[1] - 1) / stride_[1])
                pad_needed_h = (legacy_target_size_h - 1) * stride_[0] + shapes[1][2] - shapes[0][2]
                pad_needed_w = (legacy_target_size_w - 1) * stride_[1] + shapes[1][3] - shapes[0][3]
                
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pad_needed_h)/stride_[0] + 1))
                shape_out.append(int((shapes[0][3] - ((shapes[1][3] - 1)*dilation_[1] + 1) + pad_needed_w)/stride_[1] + 1))                    
        else:
            if auto_pad_mode == 'NOTSET' :
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pads_[0] + pads_[1])/stride_[0] + 1)) 
            elif auto_pad_mode == 'VALID' :
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1))/stride_[0] + 1))   
            elif auto_pad_mode == 'SAME_UPPER' or auto_pad_mode == 'SAME_LOWER' : 
                legacy_target_size_w = (shapes[0][2] + stride_[0] - 1) / stride_[0]
                pad_needed_w = (legacy_target_size_w - 1) * stride_[0] + shapes[1][2] - shapes[0][2]
                shape_out.append(int((shapes[0][2] - ((shapes[1][2] - 1)*dilation_[0] + 1) + pad_needed_w)/stride_[0] + 1))   

        return attributes, shape_out 


